import numpy as np

import torch
import torch.nn as nn
import torch.nn.utils.rnn as rnn_utils
import torch.nn.functional as F

from . import utils

HIDDEN_STATE_SIZE = 512 #RNN隐藏层的大小
EMBEDDING_DIM = 50 # 词嵌入向量的维度


class PhraseModel(nn.Module):
    def __init__(self, emb_size, dict_size, hid_size):
        '''
        语法模型

        emb_size: 词嵌入向量模型的维度
        dict_size: 单词字典（单词-> token id）的数量
        hid_size: 隐藏层单元尺寸
        '''
        super(PhraseModel, self).__init__()

        # 创建词嵌入向量层
        self.emb = nn.Embedding(num_embeddings=dict_size, embedding_dim=emb_size)
        # 创建LSTM编码器
        self.encoder = nn.LSTM(input_size=emb_size, hidden_size=hid_size,
                               num_layers=1, batch_first=True)
        # 创建LSTM解码器
        self.decoder = nn.LSTM(input_size=emb_size, hidden_size=hid_size,
                               num_layers=1, batch_first=True)
        # 创建输出层，输出维度是词token的维度，输出每个单词的概率分布
        self.output = nn.Sequential(
            nn.Linear(hid_size, dict_size)
        )

    def encode(self, x):
        '''
        将输入序列数据传入LSTM encoder层进行编码，得到输出序列以及序列最后一步的隐藏层状体
        return： 经过LSTM层得到的隐藏层结果
        '''
        '''
        当你使用 PyTorch 的 `nn.LSTM` 模块时，该模块的前向传播会返回两个主要值：`output` 和 `(h_n, c_n)`。以下是这两个返回值的详细描述：

1. **output**: 
    - Tensor 形状为 `(seq_len, batch, num_directions * hidden_size)`:
        - `seq_len`: 序列长度
        - `batch`: 批处理大小
        - `num_directions`: 如果 LSTM 是双向的，则为 2，否则为 1
        - `hidden_size`: LSTM 隐藏层的大小
    - `output` 代表了每一步的输出特征值 `h_t` 从输入层到输出层。对于双向 LSTM，它也会包括两个方向的输出。

2. **(h_n, c_n)**: 
    - `h_n`: 最后的隐藏状态。
        - Tensor 形状为 `(num_layers * num_directions, batch, hidden_size)`。
        - 对于每个层和方向，它都包含了最后时间步的隐藏状态。
    - `c_n`: 最后的 cell 状态。
        - Tensor 形状也为 `(num_layers * num_directions, batch, hidden_size)`。
        - 对于每个层和方向，它都包含了最后时间步的 cell 状态。

例如，对于一个双向 LSTM，`num_directions` 为 2。第 0 层正向的隐藏状态和 cell 状态在 `(h_n, c_n)` 的第 0 位置，第 0 层反向的隐藏状态和 cell 状态在第 1 位置，第 1 层正向的隐藏状态和 cell 状态在第 2 位置，以此类推。

简而言之，`output` 为所有时间步的输出，而 `(h_n, c_n)` 为最后一个时间步的隐藏状态和 cell 状态，它们都按层和方向排列。
        '''
        _, hid = self.encoder(x)
        return hid

    def get_encoded_item(self, encoded, index):
        # For RNN
        # 这个方法仅适用于LSTM，不适用GRU或者其他隐藏层只有单个张量的结构
        # 作用，根据index提取出对应index索引的输入对话经过编码后的隐藏层状态
        # encoded：经过encoder LSTM编码后的序列
        # index：当前是第几个对话
        #
        # return encoded[:, index:index+1]
        # For LSTM

        # 这里看似是要根据index id提取出对应当前对话的上一句输入经过编码后的隐藏层信息和最后一个时间步的cell状态
        '''
在 PyTorch 中，contiguous() 是一个经常被使用的方法，尤其当我们需要确保一个张量 (tensor) 在内存中是连续存储的时候。首先，我们需要理解什么是"连续存储"。

什么是连续存储？
在物理内存或显存中，如果一个张量的元素是顺序存储的，无间隔地紧密排列，那么我们就称这个张量为连续的。但由于某些操作（例如 transpose，permute 等），张量可能不再是连续存储的。

contiguous() 的作用：
contiguous() 方法返回一个在内存中连续存储的张量。如果原始张量已经是连续的，该方法返回原始张量的一个视图；如果原始张量不是连续的，则返回原始张量的一个连续副本。

为什么需要 contiguous()?
有些 PyTorch 函数或方法在操作前要求张量必须是连续的。例如，当你想使用 .view() 方法改变张量的形状时，这个张量必须是连续的。如果张量不是连续的，尝试调用 .view() 会引发一个错误。在这种情况下，你通常会首先调用 .contiguous()，然后再调用 .view()。
        '''
        return encoded[0][:, index:index+1].contiguous(), \
               encoded[1][:, index:index+1].contiguous()

    def decode_teacher(self, hid, input_seq):
        '''
        todo 作用

        hid： 上一层RNN网络层返回的隐藏层信息
        input_seq: 输入的对话序列，在这里输入的下一句回答的序列，格式是PackedSequence
        '''
        # Method assumes batch of size=1
        # 输入下一句回答的序列后，得到输出序列，shape是（N, L, D*H）
        out, _ = self.decoder(input_seq, hid)
        # todo 这里的data是否是input_seq输入的每一个字符的时间序列的输入对应的输出结果
        # 将结果映射输出到全连接层，应该是要得到一个输出序列每一个字符的概率概率分布 todo 作用
        out = self.output(out.data)
        return out

    def decode_one(self, hid, input_x):
        '''
        这个方法的作用就是输入一个时间序列的字符预测到下一个时间序列的字符

        hid： 上一次LSTM层输出的隐藏层信息，包含 hidden state and cell state
        input_x: 输入，一个时间序列，也就是只有一个字符

        return： (预测到的下一个时间序列的token的概率分布，隐藏层新的状态信息)
        '''
        # 输入当前的时间序列的字符 得到预测的输出和新的隐藏层信息
        out, new_hid = self.decoder(input_x.unsqueeze(0), hid)
        # 将预测的时间序列转换为单词token的概率分布
        out = self.output(out)
        return out.squeeze(dim=0), new_hid

    def decode_chain_argmax(self, hid, begin_emb, seq_len, stop_at_token=None):
        """
        Decode sequence by feeding predicted token to the net again. Act greedily

        hid: encoder LSTM层输出的隐藏层信息，包含 hidden state and cell state
        begin_emb: 对应对话的下一句回答的第一个字符的token id
        seq_len: 对应对话的下一句回答的的长度（去除了第一个token）
        stop_at_token: todo 作用 若不为None，则停止采样，采样结束后，返回的结果中只包含stop_at_token

        return： (预测到的下一个时间序列的每一个单词token的概率分布，预测到的回答token id序列)
        """
        res_logits = []
        res_tokens = []
        cur_emb = begin_emb

        for _ in range(seq_len):
            # 预测下一个时间序列的单词token的概率分布和隐藏层新的信息
            out_logits, hid = self.decode_one(hid, cur_emb)
            # 选择概率最大的得到预测到下一个时间序列的单词token id
            out_token_v = torch.max(out_logits, dim=1)[1]
            out_token = out_token_v.data.cpu().numpy()[0]

            # 将预测到的token id转换为词向量，用于下一次预测的输入时间序列字符
            cur_emb = self.emb(out_token_v)

            res_logits.append(out_logits)
            res_tokens.append(out_token)
            # 如果说有指定停止token的话，则预测到停止token的字符则停止进行预测
            if stop_at_token is not None and out_token == stop_at_token:
                break
        return torch.cat(res_logits), res_tokens

    def decode_chain_sampling(self, hid, begin_emb, seq_len, stop_at_token=None):
        """
        Decode sequence by feeding predicted token to the net again.
        Act according to probabilities

        hid: 输入对话经过编码后的隐藏层信息
        begin_emb: 对应对话的下一句回答的起始词嵌入向量
        seq_len: 对应对话的下一句回答的的长度 todo（去除了第一个token）
        stop_at_token: 结束字符token，如果包含了该token，则直接返回

        return： （预测到的每一个字符的原始概率分布序列，使用策略梯度得到的预测字符token id回答）
        """
        res_logits = [] # 预测到的每一个序列的原始概率分布
        res_actions = [] # 使用策略梯度得到的预测到的字符token id
        cur_emb = begin_emb

        # 按照长度开始预测每一个字符
        for _ in range(seq_len):
            # 得到预测的token id和新的隐藏层信息
            out_logits, hid = self.decode_one(hid, cur_emb)
            # 计算概率分布
            out_probs_v = F.softmax(out_logits, dim=1)
            out_probs = out_probs_v.data.cpu().numpy()[0]
            # 使用策略梯度，按照概率分布选择预测到的字符token id（在策略梯度中是需要执行的动作）
            action = int(np.random.choice(out_probs.shape[0], p=out_probs))
            action_v = torch.LongTensor([action]).to(begin_emb.device)
            # 将预测到的token 封装为词嵌入向量 todo 确认原始的单词字典是否是从0到n的映射
            cur_emb = self.emb(action_v)

            res_logits.append(out_logits)
            res_actions.append(action)
            if stop_at_token is not None and action == stop_at_token:
                break
        # 这里为什么要使用torch.cat
        return torch.cat(res_logits), res_actions


def pack_batch_no_out(batch, embeddings, device="cpu"):
    '''
    将输入的对话单词序列训练数据转换可训练的数据：
    即先转换有对其填充空白字符的输入矩阵，在转换为压缩后按照时间序列组合成的输入序列
    最后再转换为词向量序列

    return (词向量、压缩后的上一句对话序列)，输入对话上一句token id序列，输出对话下一句token id序列
    '''
    # 传入的batch参数必须是list
    assert isinstance(batch, list)
    # Sort descending (CuDNN requirements)
    # 按照batch中对话的长度进行逆序排序，按照原注释描述，这个应该是cuda设备要求的\
    # 这里的排序，也是因为rnn_utils.pack_padded_sequence中enforce_sorted参数决定的
    # 输入的矩阵必须是已经排过序的 书P234
    batch.sort(key=lambda s: len(s[0]), reverse=True)
    # 将对话进行解引用，将上一句对话和下一句对话分别解包到input_idx和output_idx中
    input_idx, output_idx = zip(*batch)
    # create padded matrix of inputs
    # 获取输入对话每一个对话的长度
    lens = list(map(len, input_idx))
    # 构建一个输入对话零矩阵，输入矩阵的维度是batch_size*max_len
    input_mat = np.zeros((len(batch), lens[0]), dtype=np.int64)
    # 完成后，按照顺序将输入input_idx对话内容填充到矩阵中
    for idx, x in enumerate(input_idx):
        input_mat[idx, :len(x)] = x
    input_v = torch.tensor(input_mat).to(device)
    # 这里是将矩阵进行压缩，避免传入过多的填充字符0，已解决rnn中必须要对齐
    # 不同对话导致传入过多的无效输入的问题
    # 具体函数的作用查看链接：https://zhuanlan.zhihu.com/p/374876781
    # 返回的input_seq的数据类型为PackedSequence
    input_seq = rnn_utils.pack_padded_sequence(input_v, lens, batch_first=True)
    # lookup embeddings
    # 将输入中的token id转换为词向量
    r = embeddings(input_seq.data)
    # 重新将转换为词向量的序列包装为PackedSequence类型
    emb_input_seq = rnn_utils.PackedSequence(r, input_seq.batch_sizes)
    return emb_input_seq, input_idx, output_idx


def pack_input(input_data, embeddings, device="cpu"):
    '''
    将input_data序列包装为PackedSequenceg格式
    input_data: 一句话token id序列
    '''
    input_v = torch.LongTensor([input_data]).to(device)
    # 转换为词向量序列
    r = embeddings(input_v)
    # 将该句话包装成PackedSequence类型
    # todo 为什么要这样包装，不能和上面一样使用rnn_utils.PackedSequence吗？
    return rnn_utils.pack_padded_sequence(r, [len(input_data)], batch_first=True)


def pack_batch(batch, embeddings, device="cpu"):
    '''
    将对话数据提起转换为可输入rnn网络的格式,这个方法转换为需要输入到teacher forcing模式
    return: （词向量的PackedSequence上一句话序列矩阵，词向量的PackedSequence下一句话list,上一句话的原始token id list，下一话的原始token id list）
    '''

    # 将输入对话提取转换为可训练的数据
    emb_input_seq, input_idx, output_idx = pack_batch_no_out(batch, embeddings, device)

    # prepare output sequences, with end token stripped
    # todo 验证注释
    output_seq_list = []
    # 将输出的下一句对话转换为PackedSequence类型序列
    # todo 为什么回答要采用这种方式的封装
    for out in output_idx:
        output_seq_list.append(pack_input(out[:-1], embeddings, device))
    return emb_input_seq, output_seq_list, input_idx, output_idx


def seq_bleu(model_out, ref_seq):
    '''
    从一个logits 应答序列计算BLEU得分，是在teacher foring 模式中进行计算

    model_out: 上一层LSTM层输出的时间序列结果
    ref_seq: 下一句回答的token id序列
    '''

    # 使用max得到时间序列中目前可能性最高的token id组成LSTM输出的下一句回答token id序列
    model_seq = torch.max(model_out.data, dim=1)[1]
    # 因为计算要放在cpu中，所以这里需要确保数据是在cpu中
    model_seq = model_seq.cpu().numpy()
    #
    return utils.calc_bleu(model_seq, ref_seq)
